{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# MSC Plugin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import set_env\n",
    "from pathlib import Path\n",
    "\n",
    "temp_dir = Path(\".temp\")\n",
    "temp_dir.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "参考：[【我与TVM二三事 后篇（5）】MSC之PluginBuilder](https://zhuanlan.zhihu.com/p/681450076)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import *\n",
    "from utils import _build_plugin"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": [
     "hide-output"
    ]
   },
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "Failed to build plugin under /media/pc/data/lxw/ai/tvm-book/doc/tutorials/msc/plugin/.temp/msc_plugin/install/source_torch/build, check codegen.log for detail",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mAssertionError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[0;32mIn[3], line 5\u001b[0m\n\u001b[1;32m      3\u001b[0m     frameworks\u001b[38;5;241m.\u001b[39mappend(MSCFramework\u001b[38;5;241m.\u001b[39mTENSORRT)\n\u001b[1;32m      4\u001b[0m plugin_root \u001b[38;5;241m=\u001b[39m msc_utils\u001b[38;5;241m.\u001b[39mmsc_dir(\u001b[38;5;124mf\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;132;01m{\u001b[39;00mtemp_dir\u001b[38;5;132;01m}\u001b[39;00m\u001b[38;5;124m/msc_plugin\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m----> 5\u001b[0m managers \u001b[38;5;241m=\u001b[39m _build_plugin(frameworks, plugin_root)\n\u001b[1;32m      7\u001b[0m \u001b[38;5;66;03m# test the plugin load\u001b[39;00m\n\u001b[1;32m      8\u001b[0m _test_tvm_plugin(managers[MSCFramework\u001b[38;5;241m.\u001b[39mTVM], \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mllvm\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm-book/doc/tutorials/msc/plugin/utils.py:231\u001b[0m, in \u001b[0;36m_build_plugin\u001b[0;34m(frameworks, plugin_root)\u001b[0m\n\u001b[1;32m    229\u001b[0m install_dir \u001b[38;5;241m=\u001b[39m plugin_root\u001b[38;5;241m.\u001b[39mcreate_dir(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124minstall\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m    230\u001b[0m plugin \u001b[38;5;241m=\u001b[39m _create_plugin(externs_dir)\n\u001b[0;32m--> 231\u001b[0m managers \u001b[38;5;241m=\u001b[39m build_plugins(plugin, frameworks, install_dir, externs_dir\u001b[38;5;241m=\u001b[39mexterns_dir)\n\u001b[1;32m    232\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m managers\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/contrib/msc/plugin/build.py:128\u001b[0m, in \u001b[0;36mbuild_plugins\u001b[0;34m(plugins, frameworks, workspace, codegen_config, cpp_print_config, py_print_config, externs_dir, on_debug)\u001b[0m\n\u001b[1;32m     91\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mbuild_plugins\u001b[39m(\n\u001b[1;32m     92\u001b[0m     plugins: Dict[\u001b[38;5;28mstr\u001b[39m, \u001b[38;5;28mdict\u001b[39m],\n\u001b[1;32m     93\u001b[0m     frameworks: List[\u001b[38;5;28mstr\u001b[39m],\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     99\u001b[0m     on_debug: \u001b[38;5;28mbool\u001b[39m \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mFalse\u001b[39;00m,\n\u001b[1;32m    100\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Dict[\u001b[38;5;28mstr\u001b[39m, Any]:\n\u001b[1;32m    101\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Build the plugins and load plugin manager\u001b[39;00m\n\u001b[1;32m    102\u001b[0m \n\u001b[1;32m    103\u001b[0m \u001b[38;5;124;03m    Parameters\u001b[39;00m\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m    125\u001b[0m \u001b[38;5;124;03m        The plugin managers.\u001b[39;00m\n\u001b[1;32m    126\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[0;32m--> 128\u001b[0m     codegens \u001b[38;5;241m=\u001b[39m _build_plugins(\n\u001b[1;32m    129\u001b[0m         plugins,\n\u001b[1;32m    130\u001b[0m         frameworks,\n\u001b[1;32m    131\u001b[0m         workspace,\n\u001b[1;32m    132\u001b[0m         codegen_config\u001b[38;5;241m=\u001b[39mcodegen_config,\n\u001b[1;32m    133\u001b[0m         cpp_print_config\u001b[38;5;241m=\u001b[39mcpp_print_config,\n\u001b[1;32m    134\u001b[0m         py_print_config\u001b[38;5;241m=\u001b[39mpy_print_config,\n\u001b[1;32m    135\u001b[0m         externs_dir\u001b[38;5;241m=\u001b[39mexterns_dir,\n\u001b[1;32m    136\u001b[0m         on_debug\u001b[38;5;241m=\u001b[39mon_debug,\n\u001b[1;32m    137\u001b[0m     )\n\u001b[1;32m    138\u001b[0m     managers \u001b[38;5;241m=\u001b[39m {}\n\u001b[1;32m    139\u001b[0m     \u001b[38;5;28;01mfor\u001b[39;00m name, codegen \u001b[38;5;129;01min\u001b[39;00m codegens\u001b[38;5;241m.\u001b[39mitems():\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/contrib/msc/plugin/build.py:84\u001b[0m, in \u001b[0;36m_build_plugins\u001b[0;34m(plugins, frameworks, workspace, codegen_config, cpp_print_config, py_print_config, externs_dir, on_debug)\u001b[0m\n\u001b[1;32m     73\u001b[0m codegen \u001b[38;5;241m=\u001b[39m get_codegen(\n\u001b[1;32m     74\u001b[0m     framework,\n\u001b[1;32m     75\u001b[0m     workspace,\n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     81\u001b[0m     on_debug\u001b[38;5;241m=\u001b[39mon_debug,\n\u001b[1;32m     82\u001b[0m )\n\u001b[1;32m     83\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m codegen\u001b[38;5;241m.\u001b[39mlibs_built():\n\u001b[0;32m---> 84\u001b[0m     codegen\u001b[38;5;241m.\u001b[39mbuild_libs()\n\u001b[1;32m     85\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m codegen\u001b[38;5;241m.\u001b[39mneed_manager \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m codegen\u001b[38;5;241m.\u001b[39mmanager_built():\n\u001b[1;32m     86\u001b[0m     codegen\u001b[38;5;241m.\u001b[39mbuild_manager(ops_info)\n",
      "File \u001b[0;32m/media/pc/data/lxw/ai/tvm/python/tvm/contrib/msc/plugin/codegen/codegen.py:127\u001b[0m, in \u001b[0;36mBasePluginCodeGen.build_libs\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m    125\u001b[0m             process \u001b[38;5;241m=\u001b[39m subprocess\u001b[38;5;241m.\u001b[39mPopen(command, stdout\u001b[38;5;241m=\u001b[39mlog_f, stderr\u001b[38;5;241m=\u001b[39mlog_f, shell\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n\u001b[1;32m    126\u001b[0m         process\u001b[38;5;241m.\u001b[39mwait()\n\u001b[0;32m--> 127\u001b[0m         \u001b[38;5;28;01massert\u001b[39;00m (\n\u001b[1;32m    128\u001b[0m             process\u001b[38;5;241m.\u001b[39mreturncode \u001b[38;5;241m==\u001b[39m \u001b[38;5;241m0\u001b[39m\n\u001b[1;32m    129\u001b[0m         ), \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mFailed to build plugin under \u001b[39m\u001b[38;5;132;01m{}\u001b[39;00m\u001b[38;5;124m, check codegen.log for detail\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;241m.\u001b[39mformat(\n\u001b[1;32m    130\u001b[0m             os\u001b[38;5;241m.\u001b[39mgetcwd()\n\u001b[1;32m    131\u001b[0m         )\n\u001b[1;32m    132\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_libs\u001b[38;5;241m.\u001b[39mextend([os\u001b[38;5;241m.\u001b[39mpath\u001b[38;5;241m.\u001b[39mbasename(l) \u001b[38;5;28;01mfor\u001b[39;00m l \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_lib_folder\u001b[38;5;241m.\u001b[39mlistdir()])\n\u001b[1;32m    133\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_lib_folder\u001b[38;5;241m.\u001b[39mlistdir(as_abs\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mTrue\u001b[39;00m)\n",
      "\u001b[0;31mAssertionError\u001b[0m: Failed to build plugin under /media/pc/data/lxw/ai/tvm-book/doc/tutorials/msc/plugin/.temp/msc_plugin/install/source_torch/build, check codegen.log for detail"
     ]
    }
   ],
   "source": [
    "frameworks = [MSCFramework.TORCH, MSCFramework.TVM]\n",
    "if tvm.get_global_func(\"relax.ext.tensorrt\", True) is not None:\n",
    "    frameworks.append(MSCFramework.TENSORRT)\n",
    "plugin_root = msc_utils.msc_dir(f\"{temp_dir}/msc_plugin\")\n",
    "managers = _build_plugin(frameworks, plugin_root)\n",
    "\n",
    "# test the plugin load\n",
    "_test_tvm_plugin(managers[MSCFramework.TVM], \"llvm\")\n",
    "if tvm.cuda().exist:\n",
    "    _test_tvm_plugin(managers[MSCFramework.TVM], \"cuda\")\n",
    "_test_torch_plugin(managers[MSCFramework.TORCH])\n",
    "\n",
    "# test the plugin with manager\n",
    "model_info = {\n",
    "    \"inputs\": [\n",
    "        {\"name\": \"input_0\", \"shape\": [1, 3, 224, 224], \"dtype\": \"float32\", \"layout\": \"NCHW\"}\n",
    "    ],\n",
    "    \"outputs\": [\n",
    "        {\"name\": \"output\", \"shape\": [1, 6, 218, 218], \"dtype\": \"float32\", \"layout\": \"NCHW\"}\n",
    "    ],\n",
    "    \"nodes\": {\"total\": 4, \"input\": 1, \"msc.conv2d_bias\": 1, \"MyRelu\": 1, \"nn.max_pool2d\": 1},\n",
    "}\n",
    "_test_with_manager(managers, MSCFramework.TORCH, model_info, temp_dir)\n",
    "_test_with_manager(managers, MSCFramework.TVM, model_info, temp_dir)\n",
    "if tvm.get_global_func(\"relax.ext.tensorrt\", True) is not None:\n",
    "    byoc_info = {\n",
    "        \"inputs\": [\n",
    "            {\"name\": \"input_0\", \"shape\": [1, 3, 224, 224], \"dtype\": \"float32\", \"layout\": \"NCHW\"}\n",
    "        ],\n",
    "        \"outputs\": [\n",
    "            {\"name\": \"output\", \"shape\": [1, 6, 218, 218], \"dtype\": \"float32\", \"layout\": \"\"}\n",
    "        ],\n",
    "        \"nodes\": {\"total\": 2, \"input\": 1, \"msc_tensorrt\": 1},\n",
    "    }\n",
    "    _test_with_manager(managers, MSCFramework.TENSORRT, byoc_info, temp_dir)\n",
    "\n",
    "plugin_root.destory()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xxx",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
